from torch.utils.data import Dataset
import torch
class TimeSeriesDataset(Dataset):
    def __init__(self, features, labels, time_steps):
        self.time_sequences, self.features, self.labels = self.create_sequences(features, labels, time_steps)

    def create_sequences(self, features_init, labels_init, time_steps):
        time_sequences = []
        features = []
        labels = []
        seq_num = len(labels_init) - time_steps
        for i in range(seq_num):
            time_seq = features_init[i:i + time_steps]
            time_sequences.append(time_seq)
            feature = features_init[i + time_steps]
            features.append(feature)
            label = labels_init[i + time_steps]  # take the next time-step label
            labels.append(label)
        return torch.stack(time_sequences, dim=0), torch.stack(features, dim=0), torch.stack(labels, dim=0)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.time_sequences[idx], self.features[idx], self.labels[idx]